#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project : aimix
@File    : load_file.py
@IDE     : PyCharm
@Author  : admin
@Date    : 2025/4/14 12:00
"""

import os
from moviepy.editor import VideoFileClip
from pydub.utils import mediainfo
from utils.tools import parallel_run


def load_files(file_path: str, suffix: list = None):
    """从文件夹获取文件列表"""
    video_extensions = suffix if suffix else ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv']
    video_files = []
    for file in os.listdir(file_path):
        if any(file.endswith(ext) for ext in video_extensions):
            video_files.append(os.path.join(file_path, file))

    return video_files


def get_video_info(path):
    """获取视频的时长（秒）和大小（MB）"""
    try:
        # 获取文件扩展名（小写）
        path = os.path.normpath(path)
        file_extension = path.split(".")[-1].lower()
        # 判断是音频还是视频
        if file_extension in ["mp3", "wav", "flac", "aac", "ogg"]:  # 音频文件类型
            info = mediainfo(path)
            duration = float(info["duration"])  # 时长，单位：秒
            file_size = os.path.getsize(path)  # 文件大小，单位：字节
            return {'file': path, "duration": round(duration, 2), "size": file_size}

        elif file_extension in ["mp4", "avi", "mkv", "mov", "flv"]:  # 视频文件类型
            clip = VideoFileClip(path, audio=False)  # 禁用音频加快加载
            duration = clip.duration  # 视频时长，单位：秒
            size = os.path.getsize(path) / 1024 / 1024  # 文件大小，单位：MB
            clip.close()
            return {
                'file': path,
                "duration": round(duration, 2),
                "size": f"{round(size, 2)} MB",
            }

        else:
            print("Unsupported file type.")
            return None

    except Exception as e:
        print("错误文件")
        return None


def batch_scan_folder(folder_path, suffixes=['.mp4', '.mov', '.avi', '.mkv'], max_workers=8):
    """并发读取视频信息"""
    video_paths = []
    rtv = []
    for file in os.listdir(folder_path):
        if any(file.lower().endswith(suf) for suf in suffixes):
            video_paths.append(os.path.normpath(os.path.join(folder_path, file)))
    if video_paths:
        results = parallel_run(get_video_info, video_paths, max_workers=max_workers)
        for i, r in enumerate(results, 1):
            try:
                r["id"] = i
                rtv.append(r)
            except Exception as e:
                pass
    return rtv


def load_file_count(folder_path, suffixes=['.mp4', '.mov', '.avi', '.mkv']):
    """读取文件数量"""
    video_paths = 0
    for file in os.listdir(folder_path):
        if any(file.lower().endswith(suf) for suf in suffixes):
            video_paths += 1
    return video_paths


